package com.ai.aiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

public class FileBasedChatMemory implements ChatMemory {
    private final String BADE_DIR;

    private static final Kryo kryo = new Kryo();

    static {
        kryo.setRegistrationRequired(false);
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public FileBasedChatMemory(String dir) {
        BADE_DIR = dir;
        File file = new File(dir);
        if (!file.exists()) {
            file.mkdirs();
        }

    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> conversationMessage = getOutCreateConversation(conversationId);
        conversationMessage.addAll(messages);
        saveOrCreateConversation(conversationId, conversationMessage);

    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> allMessage = getOutCreateConversation(conversationId);
        return allMessage.stream().skip(Math.max(0,allMessage.size()-lastN))
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        File conversationFile = getConversationFile(conversationId);
        if (conversationFile.exists()) {
            conversationFile.delete();
        }
    }


    private List<Message> getOutCreateConversation(String conversationId) {
        File conversationFile = getConversationFile(conversationId);
        List<Message> messages = new ArrayList<>();
        if (conversationFile.exists()) {
            try {
                Input input = new Input(new FileInputStream(conversationFile));
                while (true) {
//                    Message message = (Message) kryo.readClassAndObject(input);
//                    messages.add(message);
                    messages = kryo.readObject(input, ArrayList.class);
                    if (input.position() == input.limit()) {
                        break;
                    }
                }
                input.close();

            } catch (FileNotFoundException e) {
                throw new RuntimeException(e);
            }
        }
        return  messages;
    }

    private void saveOrCreateConversation(String conversationId, List<Message> messages) {
        File conversationFile = getConversationFile(conversationId);

        try {
            Output output= new Output(new FileOutputStream(conversationFile));
            kryo.writeObject(output,messages);
            output.close();
        } catch (FileNotFoundException e) {
            throw new RuntimeException(e);
        }

    }


    private File getConversationFile(String conversationId) {
        return new File(BADE_DIR ,conversationId + ".kryo");
    }
}
